// Script to read out of events.lhe from MadGraph and create a ROOT file events.lhe.root

using namespace std;

#include <vector>
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <cstring>
#include <TFile.h>
#include <TTree.h>
#include <TCanvas.h>
#include <TFrame.h>
#include <TH1F.h>
#include <TBenchmark.h>
#include <TRandom.h>
#include <TSystem.h>

void treeCreationLHE(char* fileName) 
{
  // Read input file and skip up to where data starts
  string line;
  const char* source = fileName;
  ifstream myfile(source);
  cout << "Reading from: " << source << endl;
  while(getline(myfile,line)) {
    if(line=="<event>") {
      break;
    }
  }

  // Open .root file
  char filename[100]="";
  const char* fileExtension = ".root";
  strcat(filename,source);
  strcat(filename,fileExtension);
  TFile *f = TFile::Open(filename,"RECREATE");   
  if (!f) { return; }
  cout << "Saving to: " << filename << endl;

  int run = 1989;

  int event = 0;
  int nparticles = 0;

  std::vector<int> vpdgId;
  std::vector<float> vpx;
  std::vector<float> vpy;
  std::vector<float> vpz;
  std::vector<float> vpT;
  std::vector<float> vp;
  std::vector<float> vE;
  std::vector<float> vM;
  std::vector<float> vphi;
  std::vector<float> veta; // pseudorapidity
  std::vector<float> vdirection; // -1 incoming, 2 intermediate, +1 final state

  // Create a TTree
  TTree *t = new TTree("tree","Tree with vectors");
  t->Branch("run",&run);
  t->Branch("event",&event);
  t->Branch("nparticles",&nparticles);
  t->Branch("vpx",&vpx);
  t->Branch("vpy",&vpy);
  t->Branch("vpz",&vpz);
  t->Branch("vpT",&vpT);
  t->Branch("vp",&vp);
  t->Branch("vE",&vE);
  t->Branch("vM",&vM);
  t->Branch("vphi",&vphi);
  t->Branch("veta",&veta);
  t->Branch("vpdgId",&vpdgId);
  t->Branch("vdirection",&vdirection);

  Int_t nevents = 0;
  while(getline(myfile,line)) {
    const float PI = atan(1)*4;
    if(line=="<event>") {
      vpx.clear();
      vpy.clear();
      vpz.clear();
      vp.clear();
      vpT.clear();
      vE.clear();
      vM.clear();
      vphi.clear();
      veta.clear();
      vpdgId.clear();
      vdirection.clear();
      nparticles = 0;
      getline(myfile,line);

      while(getline(myfile,line)) {
	if(line=="</event>") {
	  break;
	}
	else {
	  stringstream ss(line);
	  float px,py,pz,p,pT,dummy,E,M,eta,phi;
	  int pdgId, direction;
	  ss >> pdgId  >> direction >> dummy >> dummy >> dummy >> dummy >> px >> py >> pz >> E >> M;
	  p = sqrt(px*px+py*py+pz*pz);
	  pT = sqrt(px*px+py*py);
	  eta = -log(pT/p/(pz/p+1));
	  phi = atan2(py,px)+PI;
	  vpx.push_back(px);
	  vpy.push_back(py);
	  vpz.push_back(pz);
	  vpT.push_back(pT);
	  vp.push_back(p);
	  vE.push_back(E);
	  vM.push_back(M);
	  vphi.push_back(phi);
	  veta.push_back(eta);
	  vpdgId.push_back(pdgId);
	  vdirection.push_back(direction);
	  nparticles++;
	}
      }


      event = nevents;
      nevents++;
      t->Fill();
    } 
  }
  f->Write();
  myfile.close();
  cout << "Number of events read in:" << nevents << endl;
  cout << filename << " created!" << endl;
  delete f;
}

int main(int argc, char **argv) {
  std::cout << "Tree creation LHE started!" << std::endl;
  treeCreationLHE(argv[1]);
  std::cout << "Tree creation LHE done!" << std::endl;
}
